Cancel Snowflake queries when a user kills the deferred task#69635
Cancel Snowflake queries when a user kills the deferred task#69635steveahnahn wants to merge 3 commits into
Conversation
a7c612b to
a878441
Compare
0bad94d to
fbb14ac
Compare
potiuk
left a comment
There was a problem hiding this comment.
Nice change — this closes a real gap, and threading cancel_on_kill through the defer site so the trigger inherits the operator's setting is the right shape.
I verified locally: the 57 trigger + operator tests pass, the async def on_kill signature matches BaseTrigger.on_kill (the triggerer awaits it, so a sync override would have silently never run), and the positional SnowflakeSqlApiHook(conn_id, token_life_time, token_renewal_delta) args match the constructor. Building the hook inside the sync_to_async worker rather than on the event loop is the correct call. Serialization round-trip coverage and the live-verification screenshots are appreciated.
One thing to fix and two I'd like your read on — left inline. Only the asgiref one is blocking.
Drafted-by: Claude Code (Opus 4.8); reviewed by @potiuk before posting
| from collections.abc import AsyncIterator | ||
| from typing import TYPE_CHECKING, Any | ||
|
|
||
| from asgiref.sync import sync_to_async |
There was a problem hiding this comment.
This is the first and only asgiref usage in the Snowflake provider, and asgiref isn't in this provider's dependencies in providers/snowflake/pyproject.toml. It resolves today only transitively via apache-airflow, but providers should declare what they import.
Amazon and Google both do exactly this where they use sync_to_async:
"asgiref>=2.3.0; python_version < '3.14'",
"asgiref>=3.11.1; python_version >= '3.14'",Worth carrying over the >=3.11.1 floor for 3.14 rather than a bare pin.
Drafted-by: Claude Code (Opus 4.8); reviewed by @potiuk before posting
| self.token_life_time, | ||
| self.token_renewal_delta, | ||
| ) | ||
| hook.cancel_queries(self.query_ids) |
There was a problem hiding this comment.
SnowflakeSqlApiHook.cancel_queries loops _cancel_sql_api_query_execution per id with no per-query error handling, and on_kill catches only at the outermost level — so the first id that raises aborts cancellation of every id after it.
That matters most in exactly the case this PR targets: run() polls statements sequentially and never prunes self.query_ids, so on a multi-statement operator the earlier statements have typically already completed by kill time. If cancelling a completed statement raises, the still-running later statements — the ones actually burning credits — never get cancelled.
Mitigating factor: _process_response returns a dict for HTTP 422 rather than raising, so if Snowflake answers 422 for "not running" the loop survives fine. But a 404 on a reaped handle, or a 401 on token expiry, would raise.
Did your live test cover killing a multi-statement task where some statements had already finished? If 422 is what Snowflake actually returns there, this is a non-issue and worth a one-line comment saying so. If not, a per-query try/except in _cancel_queries would make it genuinely best-effort. I didn't want to assert a bug I couldn't reproduce without live credentials.
Drafted-by: Claude Code (Opus 4.8); reviewed by @potiuk before posting
| return | ||
| self.log.info("Cancelling Snowflake query ids %s", self.query_ids) | ||
| try: | ||
| await sync_to_async(self._cancel_queries)() |
There was a problem hiding this comment.
Non-blocking, more of a docs note: the triggerer wraps this in asyncio.wait_for(trigger.on_kill(), timeout=_ON_CANCEL_TIMEOUT) — [triggerer] on_kill_timeout, default 30s.
Cancellation here is N sequential blocking POSTs, each with tenacity retries on 429/503/504. A task with many statements against a slow warehouse could exceed that budget, and because sync_to_async runs in a thread the remaining cancels can't be interrupted — you'd get an on_kill() timed out warning while the thread keeps going.
Benign in practice, but the docstring's "best effort" framing might be worth extending to mention the timeout so operators know the knob exists.
Drafted-by: Claude Code (Opus 4.8); reviewed by @potiuk before posting
A deferred SnowflakeSqlApiOperator parks its running query ids in the triggerer, so the operator's own on_kill no longer runs once the task is deferred. When a user marks that task failed, clears it, or marks it success, SnowflakeSqlApiTrigger had no on_kill hook, so the Snowflake statements kept executing on the warehouse, burning compute credits, even though the operator already cancels the queries on kill in the non-deferred path. This adds on_kill to SnowflakeSqlApiTrigger to cancel the running query ids when the user acts on the deferred task, matching the behaviour already shipped for the EMR, Dataproc, BigQuery, and Dataflow triggers. The Snowflake SQL API cancel is a blocking POST with no async variant, so it runs through sync_to_async off the triggerer event loop, and the hook is built inside that worker so no connection work touches the loop. A cancel_on_kill flag on both the operator and the trigger lets users opt out.
44bf75c to
50362e6
Compare
A single failing cancel (a reaped statement handle, a transient error) would otherwise stop the remaining ids from being cancelled, leaving the still-running statements that are actually consuming credits untouched. Declare asgiref, which the trigger imports directly rather than relying on it resolving transitively through apache-airflow.
50362e6 to
687272a
Compare
Problem
SnowflakeSqlApiOperator.on_killcancels the running queries, but afterdefer()the worker is gone and it can never run.SnowflakeSqlApiTriggerhad noon_kill(base default is a no op), so killing a deferred task leaves the statements executing on the warehouse, burning credits.Change
on_killtoSnowflakeSqlApiTrigger: cancels the running query ids.POSTwith no async variant, soon_killrunscancel_queriesthroughsync_to_async, building the hook inside that worker so no connection work touches the event loop.cancel_on_killflag (defaultTrue) on operator and trigger; operator threads it into the trigger at the defer site.Live verification
deferred, trigger polling (GET /api/v2/statements/{id}returning 202).POST /api/v2/statements/{id}/canceland the statement transitioned toABORTED. Without the change it stays running.Tests
cancel_on_kill=Falseguard, flag threaded into the trigger on defer.Was generative AI tooling used to co-author this PR?
Generated-by: Claude Code (Fable 5) following the guidelines